package com.mh.utils;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.mh.common.ChatSessionManager;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.util.Arrays;
import java.util.List;

/**
 * Date:2025/2/28
 * author:zmh
 * description: AI相关代码
 **/

@Component
@Slf4j
@RequiredArgsConstructor
public class AIUtils {

    @Value("${DASHSCOPE_API_KEY}")
    public String apiKey;

    // 用户管理器
    private final ChatSessionManager sessionManager;

    /**
     * 调用AI接口
     * @return ·
     */
    public GenerationResult callWithMessage(String openId, String userContent) throws ApiException, NoApiKeyException, InputRequiredException {
        List<Message> messages = sessionManager.getMessages(openId);
        if (messages == null) {
            throw new IllegalArgumentException("无效的会话ID");
        }

        // 添加用户消息
        messages.add(createMessage(Role.USER, userContent));

        GenerationParam param = buildGenerationParam(messages);
        GenerationResult result = executeGeneration(param);

        // 添加AI回复
        messages.add(extractAssistantMessage(result));
        return result;
    }

    /**
     * 带会话管理的连续对话
     */
    public GenerationResult continuousChat(String openId, String userContent)
            throws ApiException, NoApiKeyException, InputRequiredException {

        List<Message> messages = sessionManager.getMessages(openId);
        if (messages == null) {
            throw new IllegalArgumentException("无效的会话ID");
        }

        // 添加用户消息
        messages.add(createMessage(Role.USER, userContent));

        GenerationParam param = buildGenerationParam(messages);
        GenerationResult result = executeGeneration(param);

        // 添加AI回复
        messages.add(extractAssistantMessage(result));
        return result;
    }

    private GenerationParam buildGenerationParam(List<Message> messages) {
        return GenerationParam.builder()
                .apiKey(apiKey)
                .model("qwen-plus") // 经实测qwen-plus比deepseek-v3效果更好些，至于deepseek-r1用在这就是杀鸡用牛刀了，并且在此没有推理需求，所以综合一下决定使用qwen-plus
                .messages(messages)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();
    }

    private GenerationResult executeGeneration(GenerationParam param)
            throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        return gen.call(param);
    }

    private Message createMessage(Role role, String content) {
        return Message.builder()
                .role(role.getValue())
                .content(content)
                .build();
    }

    // 取出AI回复的消息{}
    private Message extractAssistantMessage(GenerationResult result) {
        return result.getOutput().getChoices().get(0).getMessage();
    }

}
